import torch
import numpy as np

# torch.saves(state, dir) # 模型的保存
# torch.load(dir) # 模型加载


# 多线程并行化
# 主要用于数据的预处理，如图片的预加载
k = torch.get_num_threads() # 获得用于并行化CPU操作的OpenMP线程数
print(k)
# torch.set_num_threads(int) # 设定用于并行化CPU操作的OpenMP线程数

# 分布式：torch默认使用一个GPU，多个GPU情况下需要使用pytorch提供的DataParallel
# 单机多卡
# 多机多卡


# torch的设备运行支持
# 使用.to()方法将tensor在CPU、GPU之间相互移动
x = torch.tensor([1, 2, 3], dtype=torch.double, device=torch.device('cpu'))
if torch.cuda.is_available():
    device = torch.device('cuda')
    x_gpu = x.to(device) # 将cpu类型的tensor转译成gpu，同时to方法还可以更改数据类型


# torch与numpy额属于类型转换
a = np.eye(3)
a_tensor = torch.from_numpy(a)
a_np = a_tensor.numpy()
    



# torch的Variable与Autograd
# leaf：叶子张量lesf
# Variable是对Tensor的封装，操作与tensor基本一致，不同的是，每一个Variable被构建的时候，都包含三个属性：
# Variable中所包含的tensor
# tensor的梯度.grad：tensor的梯度值，每次在计算backward时都需要将前一时刻的梯度归零，否则梯度值会一直相加
# 以何种方式得到这种梯度.grad_fn：对于输入X子节点通常为None，输出结果的节点(含过程节点)的grad_fn才有效，用于指示梯度函数是哪种类型
# 
# backward函数
# torch.autograd.backward(tensors, grad_tensor=None, retain_graph=None, create_graph=False)
#   参数释义
#   tensor：用于计算梯度的tensor，torch.autograd.backward(tensors) == tensors.backward()
#   grad_tensors：计算矩阵的梯度时会用到，该参数也是一个tensor，shape一般需要和前面的tensor保持一致
#   retain_graph：通常在调用一次backward后，pytorch会自动把计算图销毁，因此若要对某个变量重复调用backward，需将该变量参数设置True
#   create_graph：为True，则创建一个专门的graph of the derivative，便于计算高阶微分
# 
# torch.autograd.grad()函数
# torch.autograd.grad(outputs, inputs, grad_outputs=None, retain_graph=None, create_graph=False, only_inputs=True, allow_unused=False)
#   参数释义
#   计算和返回outputs关于inputs的梯度和
#   outputs：因变量，既需要求导的函数；inputs：自变量，可定义多个tensor
#   only_inputs：只计算input的梯度
#   allow_unused：如为False，则计算输出出错时（因此他们的梯度永远为零）指明不使用的inputs
#
# 其他函数
# torch.autograd.enable_grad() # 启动梯度计算的上下文管理器
# torch.autograd.no_grad() # 禁止梯度计算的上下文管理器


# torch.autograd.Function()函数
#   释义
#   {每一个原始的自动求导运算实际上是两个在tensor上运行的函数
#       forward函数：根据函数定义将tensor进行前向计算，获得输出tensor
#       backward函数：返回梯度，反向传播}
#   Function函数可自定义forward函数与backward函数，最后运用apply方法执行相应计算
#   .apply()：定义在Function父类_FunctionBase中的方法
# 自定义函数示例
from torch.autograd import Function
class demo(Function): # 自定义函数
    @staticmethod # 指明静态方法，静态方法可以通过类本身直接调用，而不需要创建类的实例。
    def forward(ctx, w, x, b):
        # ctx：上下文管理器的定义
        ctx.save_for_backward(w, x, b) # 用于存储涉及的变量，用于backward
        return w*x + b

    @staticmethod
    def backward(ctx, grad_outputs): # type: ignore
        # grad_outputs：上一级的梯度
        w, x, b = ctx.saved_tensors
        grad_w = grad_outputs * x
        grad_x = grad_outputs * w
        grad_b = grad_outputs
        return grad_w, grad_x, grad_b   
w = torch.rand(2, 2, requires_grad=True)
x = torch.rand(2, 2, requires_grad=True)
b = torch.rand(2, 2, requires_grad=True)
y_out:torch.Tensor = demo.apply(w, x, b) # type: ignore
y_out.backward(torch.ones(2, 2), retain_graph=True)
print(w, x, b)
print(w.grad, x.grad, b.grad)